#pragma once

#include <iostream>
#include <opencv2/opencv.hpp>
#include<opencv2/ml.hpp>
#include "../General/General.h"
#include "YOLOX/yolox.h"
#define USE_SVM
#define SHOW_RESULT
//#define GET_ARMOR_PIC

namespace rm
{
static const int W = 64;
static const int H = 32;

enum ArmorFlag
{
    NO,
    WHOLE,
    HALF
};

struct Robot
{
    int index;    //index in objects
    int number;
    std::vector<int> armors_of_robot; //index of armors
    cv::Rect rect;
};

struct Armor
{
    int index;
    int color; //green for death
    int number;
    int robot;  //index of robot
    int type;   //big or small
    int flag;   //whole or half
    cv::Rect rect;
    std::vector<cv::Point2f> vertex;  //vertex for PNP (tl, tr, br, bl)
    cv::Point2f center;
    float score;
    float delta_x; // 装甲板中心到机器人中心在x方向上的距离
    Armor()
    {
        index = -1;
        color = -1;
        number = -1;
        robot = -1;
        type = UNKNOWN_ARMOR;
        flag = NO;
        vertex.clear();
        center = cv::Point2f(0, 0);
        delta_x = -1;
        score = 0;
    }
};

class FourPointsDetector
{
public:
    FourPointsDetector()
    {
        vertex.clear();
        armor_num = -1;
        type = UNKNOWN_ARMOR;
        frontImg = cv::Mat();
        center = cv::Point2f(0, 0);
        armor_small_armor_ratio = 2.35;
        armor_big_armor_ratio = 3;
        brightness_threshold = 160;
        score = 0;
        _allCnt = 0;
    }
    void clear()
    {
        vertex.clear();
        armor_num = -1;
        type = UNKNOWN_ARMOR;
        frontImg = cv::Mat();
        center = cv::Point2f(0, 0);
        score = 0;
    }
    ~FourPointsDetector(){}
    void init();
    int detect(cv::Mat image);  //find lightbar vertex
    void getFrontImg(const cv::Mat& grayImg);
    int isArmorPattern();     //get armor number
    std::vector<cv::Point2f> vertex;  //vertex for PNP (tl, tr, br, bl)
    int armor_num;
    int type;
    cv::Point2f center;
    float score;

private:
    std::vector<cv::Point2f> _vertex;  //vertex for SVM
    cv::Mat frontImg;
    //4points CNN
    cv::dnn::Net net;
    //评分体系
    float armor_small_armor_ratio;
    float armor_big_armor_ratio;
    int brightness_threshold;
    //
    int _allCnt;
};

class EnemyDetector
{
public:
    EnemyDetector() {}
    void init();
    void setEnemyColor(int _enemy_color)
    {
        enemy_color = _enemy_color;
    }
    void loadImg(const cv::Mat& _srcImg);
    bool detect();     //target find -> true
    const std::vector<cv::Point2f> getTargetVertex() const;
    cv::Point2f getTargetCenter();
    int getArmorNum();
    int getArmorType();
    int getArmorFlag();
    cv::Mat getDebugImg();

private:
    int enemy_color;
    Yolox yolox_dectector;
    FourPointsDetector p4;
    cv::Mat srcImg;
    std::vector<Object> objects;
    std::vector<Robot> robots;
    std::vector<Armor> armors;

    cv::Point2f target_center;
    std::vector<cv::Point2f> target_vertex;
    int target_type;
    int target_num;
    int target_flag;
};

}
